from z3 import *


def list_to_string(expr, no_bracket=False, workaround=False):
    if type(expr) == str:
        return expr
    if type(expr) == tuple:
        return str(expr[1])
    sub_expr = []
    for expr in expr:
        if type(expr) == list:
            sub_expr.append(list_to_string(expr))
        elif type(expr) == tuple:
            sub_expr.append(str(expr[1]))
        else:
            sub_expr.append(expr)
    if no_bracket:
        return "%s" % (' '.join(sub_expr))
    if workaround and len(sub_expr) == 1:
        return "%s" % (' '.join(sub_expr))
    return "(%s)" % (' '.join(sub_expr))


def get_var_def(x):
    if x[1] == 'Int':
        return Int(x[0])
    elif x[1] == 'Bool':
        return Bool(x[0])
    elif x[1] == 'Real':
        return Real(x[0])
    else:
        raise ValueError("Unsupported type: %s" % str(x[1]))


class Checker:
    def __init__(self, bm_expr):
        constraint = []
        self.var_def = {}
        for expr in bm_expr:
            if 0 == len(expr):
                continue
            elif expr[0] == 'constraint':
                constraint.append('(assert %s)' % list_to_string(expr[1:], workaround=True))
            elif expr[0] == 'declare-var':
                self.var_def[expr[1]] = get_var_def(expr[1:])
        self.constraint_str = '\n'.join(constraint)
        self.solver = Solver()

    def get_counterexample(self, fun_def):
        smt2 = fun_def + '\n' + self.constraint_str
        cx = parse_smt2_string(smt2, decls=self.var_def)
        cx = And(cx)
        self.solver.push()
        self.solver.add(Not(cx))
        if self.solver.check() == unsat:
            self.solver.pop()
            return None
        else:
            # result = self.solver.model()
            self.solver.pop()
            # return result
            return []  # we don't use counterexample now
